# neo4j_utils/connect.py
from neo4j import GraphDatabase

class Neo4jConnection:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        if self.driver:
            self.driver.close()

    def create_node(self, label, properties):
        with self.driver.session() as session:
            session.execute_write(self._create_node_tx, label, properties)

    @staticmethod
    def _create_node_tx(tx, label, properties):
        # 使用MERGE替代CREATE避免重复创建相同name的节点
        query = f"""
        MERGE (n:{label} {{name: $properties.name}})
        SET n += $properties
        RETURN n
        """
        tx.run(query, properties=properties)

    def create_relationship(self, start_node_label, end_node_label, rel_type, properties=None):
        with self.driver.session() as session:
            session.execute_write(
                self._create_relationship_tx,
                start_node_label,
                end_node_label,
                rel_type,
                properties or {}
            )

    # neo4j_utils/connect.py（修改后）
    @staticmethod
    def _create_relationship_tx(tx, start_label, end_label, rel_type, properties):
        # 根据节点类型动态匹配属性（国家-首都用capital，国家-地球村用village）
        end_property = "capital" if end_label == "Capital" else "village"
        query = f"""
        MATCH (a:{start_label} {{name: $country}}), (b:{end_label} {{name: ${end_property}}})
        MERGE (a)-[r:{rel_type}]->(b)
        SET r += $props
        """
        tx.run(
            query,
            country=properties['country'],
            **{end_property: properties[end_property]},  # 动态传递end属性
            props=properties
        )